将 Jupyter Notebook 作为模块导入

人们普遍遇到的一个问题是希望从 Jupyter Notebook 中导入代码。但由于 Notebook 不是纯 Python 文件，因此无法通过常规的 Python 机制导入，这使得问题变得困难。

幸运的是，Python 为导入机制提供了一些相当复杂的钩子（hooks），所以我们实际上可以毫不费力地让 Jupyter Notebook 变得可导入，并且只使用公开的 API。

import io, os, sys, types

from IPython import get_ipython
from nbformat import read
from IPython.core.interactiveshell import InteractiveShell

导入钩子通常采用两种对象的形式：

1. 一个模块**加载器**（Loader），它接收一个模块名（例如 'IPython.display'），并返回一个模块（Module）。

2. 一个模块**查找器**（Finder），它判断一个模块是否存在，并告诉 Python 使用哪个**加载器**（Loader）。

def find_notebook(fullname, path=None):
    """find a notebook, given its fully qualified name and an optional path

    This turns "foo.bar" into "foo/bar.ipynb"
    and tries turning "Foo_Bar" into "Foo Bar" if Foo_Bar
    does not exist.
    """
    name = fullname.rsplit('.', 1)[-1]
    if not path:
        path = ['']
    for d in path:
        nb_path = os.path.join(d, name + ".ipynb")
        if os.path.isfile(nb_path):
            return nb_path
        # let import Notebook_Name find "Notebook Name.ipynb"
        nb_path = nb_path.replace("_", " ")
        if os.path.isfile(nb_path):
            return nb_path

Notebook 加载器

这里是我们的 Notebook 加载器。它实际上非常简单——一旦我们确定了模块的文件名，它所做的就是：

1. 将 notebook 文档加载到内存中

2. 创建一个空的模块（Module）

3. 在模块（Module）的命名空间中执行每个单元格

由于 IPython 单元格可以包含扩展语法，因此在执行之前会应用 IPython 转换，将这些单元格转换为纯 Python 对应物。如果你所有的 notebook 单元格都是纯 Python 代码，则此步骤是不必要的。

class NotebookLoader(object):
    """Module Loader for Jupyter Notebooks"""

    def __init__(self, path=None):
        self.shell = InteractiveShell.instance()
        self.path = path

    def load_module(self, fullname):
        """import a notebook as a module"""
        path = find_notebook(fullname, self.path)

        print("importing Jupyter notebook from %s" % path)

        # load the notebook object
        with io.open(path, 'r', encoding='utf-8') as f:
            nb = read(f, 4)

        # create the module and add it to sys.modules
        # if name in sys.modules:
        #    return sys.modules[name]
        mod = types.ModuleType(fullname)
        mod.__file__ = path
        mod.__loader__ = self
        mod.__dict__['get_ipython'] = get_ipython
        sys.modules[fullname] = mod

        # extra work to ensure that magics that would affect the user_ns
        # actually affect the notebook module's ns
        save_user_ns = self.shell.user_ns
        self.shell.user_ns = mod.__dict__

        try:
            for cell in nb.cells:
                if cell.cell_type == 'code':
                    # transform the input to executable Python
                    code = self.shell.input_transformer_manager.transform_cell(cell.source)
                    # run the code in themodule
                    exec(code, mod.__dict__)
        finally:
            self.shell.user_ns = save_user_ns
        return mod

模块查找器

查找器是一个简单的对象，它告诉你一个名称是否可以被导入，并返回相应的加载器。它所做的只是检查，当你执行

import mynotebook

它会检查 mynotebook.ipynb 是否存在。如果找到了 notebook，那么它会返回一个 NotebookLoader。

任何额外的逻辑都只是为了解析包（package）内的路径。

class NotebookFinder(object):
    """Module finder that locates Jupyter Notebooks"""

    def __init__(self):
        self.loaders = {}

    def find_module(self, fullname, path=None):
        nb_path = find_notebook(fullname, path)
        if not nb_path:
            return

        key = path
        if path:
            # lists aren't hashable
            key = os.path.sep.join(path)

        if key not in self.loaders:
            self.loaders[key] = NotebookLoader(path)
        return self.loaders[key]

注册钩子

现在我们将 NotebookFinder 注册到 sys.meta_path 中

sys.meta_path.append(NotebookFinder())

从现在开始，我的 notebook 应该就可以被导入了。

让我们看看当前工作目录（CWD）里有什么

ls nbpackage

所以我应该能够 import nbpackage.mynotebook。

import nbpackage.mynotebook

题外话：显示 notebook

这里有一些简单的代码，用于显示带有语法高亮等的 notebook 内容。

from pygments import highlight
from pygments.lexers import PythonLexer
from pygments.formatters import HtmlFormatter

from IPython.display import display, HTML

formatter = HtmlFormatter()
lexer = PythonLexer()

# publish the CSS for pygments highlighting
display(
    HTML(
        """
<style type='text/css'>
%s
</style>
"""
        % formatter.get_style_defs()
    )
)

def show_notebook(fname):
    """display a short summary of the cells of a notebook"""
    with io.open(fname, 'r', encoding='utf-8') as f:
        nb = read(f, 4)
    html = []
    for cell in nb.cells:
        html.append("<h4>%s cell</h4>" % cell.cell_type)
        if cell.cell_type == 'code':
            html.append(highlight(cell.source, lexer, formatter))
        else:
            html.append("<pre>%s</pre>" % cell.source)
    display(HTML('\n'.join(html)))


show_notebook(os.path.join("nbpackage", "mynotebook.ipynb"))

所以我的 notebook 有一些代码单元格，其中一个包含了一些 IPython 语法。

让我们看看导入它时会发生什么

from nbpackage import mynotebook

太棒了，它导入成功了！它能工作吗？

mynotebook.foo()

再次欢呼！

即使是包含 IPython 语法的函数也能工作

mynotebook.has_ip_syntax()

包中的 Notebook

我们在 nb 包里也有一个 notebook，所以让我们确保它也能正常工作。

ls nbpackage/nbs

请注意，__init__.py 文件是必需的，这样 nb 才能被视为一个包，这和通常情况一样。

show_notebook(os.path.join("nbpackage", "nbs", "other.ipynb"))

from nbpackage.nbs import other

other.bar(5)

所以现在我们可以从本地目录和包内部导入 notebook 了。

我甚至可以在 IPython 中放入一个 notebook，以进一步证明这能正常工作

import shutil
from IPython.paths import get_ipython_package_dir

utils = os.path.join(get_ipython_package_dir(), 'utils')
shutil.copy(
    os.path.join("nbpackage", "mynotebook.ipynb"), os.path.join(utils, "inside_ipython.ipynb")
)

并从 IPython.utils 导入该 notebook

from IPython.utils import inside_ipython

inside_ipython.whatsmyname()

这种方法甚至可以导入在 notebook 中使用 %%cython 魔法命令定义的函数和类。